import re
import json
from py2neo import Graph
from flask import Flask, request, jsonify, make_response
from flask_cors import CORS
import zhipuai
from sentence_transformers import SentenceTransformer
import time

# -------------------------- 数据库连接与查询函数 --------------------------
try:
    graph = Graph("neo4j://localhost:7687", auth=("neo4j", "053116wj"))
    print("成功连接到图数据库！")
except Exception as e:
    print(f"连接数据库失败：{e}")
    graph = None


def execute_query_and_get_data(query: str) -> list:
    if not graph:
        return []
    try:
        return graph.run(query).data()
    except Exception as e:
        print(f"查询出错: {e}")
        return []


# -------------------------- AI模型调用 --------------------------
zhipuai.api_key = "743c5a9da68643e29d05d116288811ae.zUHCY7Jd1tniSCPV"  # 请替换为您的智谱AI API密钥
client = zhipuai.ZhipuAI(api_key=zhipuai.api_key)

# 加载用于向量检索的嵌入模型
embedding_model = SentenceTransformer('moka-ai/m3e-base')


def call_glm_4_5(user_question: str, db_result: str) -> str:
    # 构建提示词，如果数据库有结果，则让AI基于结果进行回答
    # 否则，让AI自由回答
    if db_result:
        prompt = (
            f"""
            角色：你是一个卷烟行业增加卷烟甜味的专家，擅长回答提出的有关卷烟甜味的问题。这些问题主要包含致甜成分的各种属性、实验方式、研究论文等
            用户的问题是：{user_question}
            知识图谱信息：{db_result}
            请根据以上信息进行回答。
            回答要求：
            1、请你综合以上信息，用接近人类的交流方式，流畅、准确地回答用户的问题。
            2、特别注意：主要参考查询返回的结果，再结合自己预训练学习的知识进行回答，不要直接输出查询返回的结果！！！
            3、回答的语气尽量保持学术性，多总结要点。
            4、特别注意：回答的结果不要带有Markdown格式的特殊符号以及任何其他的符号！！！
            5、回答的开头为：结合专业知识。
            6、回答的内容综合查询返回结果和预训练自身具备的知识。
            7、回答的内容要包含查询到的具体实体名称，具体的论文名称等。
            8、对于添加量的查询，给出具体的添加数值或者添加量的范围，注意要给出相应的单位。
            9、回答的内容需要给出相应的依据，依据为查询到的内容，最好为具体的论文名称，每个回答内容给出3-5条依据，给出的依据一定要是知识图谱存储的论文名称。

            """
        )
    else:
        prompt = (
            f"""
            角色：你是一个卷烟行业增加卷烟甜味的专家，擅长回答提出的有关卷烟甜味的问题。这些问题主要包含致甜成分的各种属性、实验方式、研究论文等
            用户的问题是：{user_question}
            请直接给出回答。
            回答要求：
            1、结合自身的预训练知识回答用户提出的问题
            2、回答的开头为：我暂未查询到相关知识，但据我了解
            3、回答的结尾为：请您继续查询相关资料，确保信息的真实性！
            """)

    try:
        response = client.chat.completions.create(
            model="glm-4.5",  # 这里替换为可用的模型
            messages=[
                {"role": "user", "content": prompt}
            ],
            stream=False,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"AI模型调用失败: {e}"


# -------------------------- 检索函数 --------------------------

def generate_cypher_query_with_llm(user_question: str) -> str:
    """
    使用LLM根据用户问题生成Cypher查询。
    """
    # 知识图谱的节点和关系信息，作为给LLM的上下文
    graph_schema = """
    重要提示：请严格按照以下提供的完整节点标签和关系类型进行查询，不要进行任何简化或省略。例如，致甜成分的节点标签是“致甜成分名称”，而不是“致甜成分”。
    图数据库中包含以下节点标签(Node Labels)和关系类型(Relationship Types):
    - 节点标签:
      - 致甜成分名称: 表示卷烟中的致甜成分。例如：麦芽糖
      - 甜味剂类型: 表示甜味剂的分类，例如“天然植物提取物”，“合成甜味剂”
      - 论文名称: 表示相关的研究论文。
      - 致甜成分的香气: 表示致甜成分的香气特征。
      - 致甜成分的优点: 表示致甜成分的优点。
      - 致甜成分的缺点: 表示致甜成分的缺点。
      - 致甜成分的应用: 表示致甜成分的应用场景。
      - 致甜成分添加量及口感: 表示致甜成分的添加量及口感。
      - 致甜成分实验方法: 表示致甜成分的实验或制备方法。
      - 致甜成分的添加方法: 表示致甜成分的添加方法。
      - 致甜成分相对甜度: 表示致甜成分的相对甜度。
      - 甜味持久性: 表示致甜成分甜味的持久性。
      - 致甜成分烟气烟丝中的含量及迁移率: 表示致甜成分在烟丝烟气中的含量和迁移率。
      - 烟叶类型: 表示与致甜成分相关的烟叶类型。
      - 评吸方法: 表示致甜成分的评吸方法。

    - 关系类型:
      - 香气
      - 研究论文
      - 致甜物质
      - 优点
      - 缺点
      - 应用于
      - 添加量及口感
      - 实验方法
      - 添加方式
      - 相对甜度
      - 持久性
      - 含量及迁移率
      - 相关论文
      - 分析方法
      - 烟叶类别
      - 评吸方式

    任务：根据用户的问题，利用上述知识图谱信息，生成一个**有效的Cypher查询语句**。
    1. 你的回答必须是纯粹的Cypher查询语句，不要包含任何解释性文字或Markdown格式。
    2. 查询语句必须是单行，且结尾必须是有效的RETURN子句。
    3. 对于具体的致甜成分名称，例如：麦芽酚，需要查询与致甜成分名称相关的所有关系和实体，即通过添加方式、余味、分析方法、应用于等关系编写Cypher语句进行查询。
    4. 如果问题无法用Cypher查询，请调整相应关系多次尝试。
    5. 所有节点的主键属性都是 'name'。
    6. 返回的目标变量（如e或m）必须使用 RETURN <变量>.name AS target 的格式。
    7. 若用户输入的问题只包含实体名称，则用所有与实体相关的关系进行查询。
    8. 生成的查询语句必须参考给出的可用查询语句
    
    可用查询语句：
    查找香气：
    MATCH (:致甜成分名称 {name: '麦芽糖'})-[:香气]->(m:致甜成分的香气) RETURN m.name AS target

    查找研究论文：
    MATCH (:致甜成分名称 {name: '蔗糖'})<-[:致甜物质]-(m:论文名称) RETURN m.name AS target

    查找优点：
    MATCH (:致甜成分名称 {name: '甜菊糖苷'})-[:优点]->(m:致甜成分的优点) RETURN m.name AS target

    查找缺点：
    MATCH (:致甜成分名称 {name: '阿斯巴甜'})-[:缺点]->(m:致甜成分的缺点) RETURN m.name AS target

    查找应用：
    MATCH (:致甜成分名称 {name: '木糖醇'})-[:应用于]->(m:致甜成分的应用) RETURN m.name AS target

    查找添加量及口感：
    MATCH (:致甜成分名称 {name: '麦芽糖醇'})-[:添加量及口感]->(m:致甜成分添加量及口感) RETURN m.name AS target

    查找实验方法：
    MATCH (:致甜成分名称 {name: '乳糖'})-[:实验方法]->(m:致甜成分实验方法) RETURN m.name AS target

    查找添加方式：
    MATCH (:致甜成分名称 {name: '蔗糖'})-[:添加方式]->(m:致甜成分的添加方法) RETURN m.name AS target

    查找相对甜度：
    MATCH (:致甜成分名称 {name: '甜菊糖苷'})-[:相对甜度]->(m:致甜成分相对甜度) RETURN m.name AS target

    查找持久性：
    MATCH (:致甜成分名称 {name: '三氯蔗糖'})-[:持久性]->(m:甜味持久性) RETURN m.name AS target

    查找含量及迁移率：
    MATCH (:致甜成分名称 {name: '麦芽酚'})-[:含量及迁移率]->(m:致甜成分烟气烟丝中的含量及迁移率) RETURN m.name AS target

    查找烟叶类型：
    MATCH (:致甜成分名称 {name: '阿斯巴甜'})-[:烟叶类别]->(m:烟叶类型) RETURN m.name AS target

    查找评吸方法：
    MATCH (:致甜成分名称 {name: '蔗糖'})-[:评吸方式]->(m:评吸方法) RETURN m.name AS target

    从论文查找致甜成分：
    MATCH (:论文名称 {name: '一种甜味剂在烟草中的应用'})-[:致甜物质]->(e:致甜成分名称) RETURN e.name AS target

    从甜味剂类型查找论文：
    MATCH (:甜味剂类型 {name: '天然植物提取物'})-[:研究论文]->(e:论文名称) RETURN e.name AS target

    查找所有致甜成分：
    MATCH (e:致甜成分名称) RETURN e.name AS target

    查找所有甜味剂类型：
    MATCH (e:甜味剂类型) RETURN e.name AS target

    查找所有论文：
    MATCH (e:论文名称) RETURN e.name AS target


    - 示例：
        - 用户问题: "《一种甜味剂在烟草中的应用》中提到了哪些致甜成分？"
        - 你的回答: "MATCH (:论文名称 {{name: '一种甜味剂在烟草中的应用'}})-[:致甜物质]->(e:致甜成分名称) RETURN e.name AS target"
        - 用户问题: "蔗糖的甜度是多少？"
        - 你的回答: "MATCH (:致甜成分名称 {{name: '蔗糖'}})-[:相对甜度]->(m:致甜成分相对甜度) RETURN m.name AS target"
    """

    messages = [
        {"role": "system", "content": graph_schema},
        {"role": "user", "content": f"请为以下用户问题生成Cypher查询：\n\n用户问题: {user_question}"}
    ]

    # ------------------ 调用AI生成Cypher ------------------
    try:
        response = client.chat.completions.create(
            model="glm-4.5",
            messages=[
                {"role": "system", "content": "你是一个将中文问题转化为Cypher查询的助手，根据给定的图谱模式进行转换。"},
                {"role": "user", "content": graph_schema + "\n用户问题：" + user_question}
            ],
            stream=False,
        )
        cypher_query = response.choices[0].message.content.strip().replace("`", "").replace("cypher", "")
        # 简单校验查询是否有效
        if "MATCH" in cypher_query and "RETURN" in cypher_query:
            return cypher_query
        return "NOT_APPLICABLE"
    except Exception as e:
        print(f"AI生成Cypher失败: {e}")
        return "NOT_APPLICABLE"



def get_vector_search_result(question: str) -> list:
    """
    使用 GDS 函数加速向量检索，然后执行图结构混合检索。
    """

    # 确保 embedding_model, graph, question 已定义
    if not graph:
        print("向量检索失败：数据库连接对象不存在。")
        return []

    try:
        # -------------------------- 阶段零：Python 逻辑 --------------------------
        # 编码用户问题，并转换为 Cypher 参数所需的 List 格式
        encoded_vectors = embedding_model.encode([question], normalize_embeddings=False)
        query_embedding = encoded_vectors[0].tolist()

        params = {
            "queryEmbedding": query_embedding,
            "min_similarity": 0.8  # 您设置的相似度阈值
        }

        # -------------------------- Cypher 查询（GDS 加速修正版） --------------------------
        cypher_query = """
        // 阶段一：GDS 向量加速检索 - 锚定语义最相关的实体
        MATCH (n:致甜成分名称)
        WHERE n.embedding IS NOT NULL

        // 1. 使用 GDS.similarity.cosine 替代 REDUCE 手动计算
        WITH n, gds.similarity.cosine(n.embedding, $queryEmbedding) AS similarity

        // 2. 筛选并锚定得分最高的 5 个实体
        WHERE similarity > $min_similarity
        WITH n, similarity
        ORDER BY similarity DESC LIMIT 5

        // 阶段二：图结构查询 - 获取锚定实体的全部邻居信息和关系
        MATCH (n)-[r]-(neighbor)

        // 3. 构造并收集所有结果 (解决归属和作用域问题)
        WITH n, similarity,
            // 锚定节点自身
            {target: n.name, type: labels(n)[0], similarity: similarity} AS anchor_result,
            // 邻居节点
            collect(DISTINCT {target: neighbor.name, type: labels(neighbor)[0], similarity: similarity}) AS neighbor_results,
            // 关系
            collect(DISTINCT {target: TYPE(r) + " (ID: " + ID(r) + ")", type: TYPE(r), similarity: similarity}) AS relation_results

        // 4. 合并所有列表
        WITH [anchor_result] + neighbor_results + relation_results AS combined_results

        // 5. UNWIND 列表为最终的行结果
        UNWIND combined_results AS row

        // 6. 最终返回和排序
        RETURN DISTINCT row.target AS target, row.type AS type, row.similarity AS similarity
        ORDER BY similarity DESC
        LIMIT 20
        """

        # 运行查询时，传递参数
        return graph.run(cypher_query, **params).data()

    except Exception as e:
        # 如果 GDS 在生产环境中遇到新问题，至少能捕获错误
        print(f"向量检索失败: {e}")
        return []

# -------------------------- 主问答函数 (已优化) --------------------------
def answer_question(question: str) -> str:
    if not graph:
        return "抱歉，无法连接到图数据库。"

    # 第一步：尝试让LLM生成Cypher查询
    cypher_query = generate_cypher_query_with_llm(question)
    result_data = []

    if cypher_query != "NOT_APPLICABLE":
        print(f"✅ LLM成功生成了Cypher查询: {cypher_query}")
        result_data = execute_query_and_get_data(cypher_query)

    # 如果第一步没有得到结果，则降级到向量检索
    if not result_data:
        print("🔍 Cypher查询失败或无结果，正在进行向量检索...")
        result_data = get_vector_search_result(question)

    db_result = ""
    if result_data:
        # ⚠️ 关键调试修改：打印原始的字典列表，而不是拼接后的字符串
        # 这样就能清晰地看到哪个知识点（d['target']）包含了具体的数字
        print(f"🔍 原始检索结果字典 (result_data):")
        # 使用 pprint 或 json.dumps 格式化打印，提高可读性
        import json
        print(json.dumps(result_data, indent=4, ensure_ascii=False))

        # 组合所有结果的name属性，作为上下文（这行保持不变，给LLM提供上下文）
        db_result = " ".join([str(d['target']) for d in result_data])
        print(f"✅ 检索到数据库结果 (db_result for LLM): {db_result}")
    else:
        print("❌ 未从数据库中检索到相关信息。")

    # 最后，将问题和结果交给AI进行最终回答
    return call_glm_4_5(question, db_result)


# -------------------------- 新增：Flask后端接口 --------------------------
app = Flask(__name__)
CORS(app)  # 解决前端跨域问题（重要，否则前端无法调用接口）

# 接口1：处理前端的查询请求（接收问题，返回答案）
@app.route('/api/query', methods=['POST'])
def handle_query():
    # 接收前端发送的JSON数据（包含问题和用户信息，此处简化）
    data = request.get_json()
    question = data.get('question', '').strip()  # 前端传来的问题
    if not question:
        return jsonify({"answer": "请输入有效的查询问题！"})
    # 调用原有问答函数获取结果
    answer = answer_question(question)
    # 返回结果给前端
    return jsonify({"answer": answer})

# 接口2：模拟登录（简化版，实际项目需加密码验证）
@app.route('/api/login', methods=['POST'])
def handle_login():
    data = request.get_json()
    username = data.get('username', '')
    password = data.get('password', '')
    # 此处简化：仅验证非空（实际项目需对接用户数据库，验证账号密码）
    if username and password:
        return jsonify({"success": True, "message": "登录成功"})
    else:
        return jsonify({"success": False, "message": "账号或密码不能为空"})

if __name__ == '__main__':
    # 启动后端服务（默认端口5000，前端需指向这个地址）
    app.run(host='0.0.0.0', port=5000, debug=True)